package ynu.edu.service.impl;

import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.ByteArrayResource;
import org.springframework.http.*;
import org.springframework.stereotype.Service;
import org.springframework.util.LinkedMultiValueMap;
import org.springframework.util.MultiValueMap;
import org.springframework.web.client.RestTemplate;
import ynu.edu.domain.PageDTO;
import ynu.edu.domain.dto.Img2ImgRequest;
import ynu.edu.domain.vo.Img2ImgResponse;
import ynu.edu.entity.Img2Img;
import ynu.edu.mapper.Img2ImgMapper;
import ynu.edu.service.Img2ImgService;
import ynu.edu.utils.AliOssUtil;

import java.time.LocalDateTime;
import java.util.*;

/**
 * Img2Img服务实现类
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class Img2ImgServiceImpl extends ServiceImpl<Img2ImgMapper, Img2Img> implements Img2ImgService {

    private final RestTemplate restTemplate;
    private final AliOssUtil aliOssUtil;
    private final ObjectMapper objectMapper;

    // Flask API地址
    @Value("${flask.api.url:http://127.0.0.1:5000}")
    private String flaskApiUrl;

    @Override
    public Img2ImgResponse img2img(Img2ImgRequest request) {
        // 根据modelName决定调用哪个风格转换
        String modelName = request.getModelName();
        
        if (modelName != null) {
            if (modelName.toLowerCase().contains("ghibli")) {
                return ghibliStyleTransform(request);
            } else {
                return animeStyleTransform(request);
            }
        } else {
            // 默认使用动漫风格
            return animeStyleTransform(request);
        }
    }

    @Override
    public Img2ImgResponse animeStyleTransform(Img2ImgRequest request) {
        try {
            // 1. 提取Base64图像数据
            String base64Image = extractBase64Data(request.getInitImage());
            
            // 2. 调用Flask API生成图片
            Map<String, Object> response = callFlaskApi("/api/anime-style", base64Image, request.getPrompt(), request.getNegativePrompt());
            
            // 3. 处理响应并生成结果
            return processApiResponse(response, request, "anime_style");
            
        } catch (Exception e) {
            log.error("动漫风格转换失败", e);
            throw new RuntimeException("动漫风格转换失败: " + e.getMessage(), e);
        }
    }

    @Override
    public Img2ImgResponse ghibliStyleTransform(Img2ImgRequest request) {
        log.info("开始宫崎骏风格转换: 提示词：{}，负面提示词：{}，图片的长度：{}",
                request.getPrompt(), request.getNegativePrompt(), request.getInitImage().length());
        
        try {
            // 1. 提取Base64图像数据
            String base64Image = extractBase64Data(request.getInitImage());
            
            // 2. 调用Flask API生成图片
            Map<String, Object> response = callFlaskApi("/api/ghibli-style", base64Image, request.getPrompt(), request.getNegativePrompt());
            
            // 3. 处理响应并生成结果
            return processApiResponse(response, request, "ghibli_style");
            
        } catch (Exception e) {
            log.error("宫崎骏风格转换失败", e);
            throw new RuntimeException("宫崎骏风格转换失败: " + e.getMessage(), e);
        }
    }

    @Override
    public Map<String, Object> checkServiceHealth() {
        try {
            String url = flaskApiUrl + "/api/health";
            ResponseEntity<Map> response = restTemplate.getForEntity(url, Map.class);
            
            if (response.getStatusCode() == HttpStatus.OK && response.getBody() != null) {
                return response.getBody();
            } else {
                Map<String, Object> errorResult = new HashMap<>();
                errorResult.put("status", "error");
                errorResult.put("message", "服务健康检查失败，状态码：" + response.getStatusCode());
                return errorResult;
            }
        } catch (Exception e) {
            log.error("健康检查失败", e);
            Map<String, Object> errorResult = new HashMap<>();
            errorResult.put("status", "error");
            errorResult.put("message", "健康检查失败: " + e.getMessage());
            return errorResult;
        }
    }

    @Override
    public Img2Img getByTaskId(String taskId) {
        LambdaQueryWrapper<Img2Img> wrapper = new LambdaQueryWrapper<>();
        wrapper.eq(Img2Img::getTaskId, taskId);
        return getOne(wrapper);
    }

    @Override
    public List<Img2Img> getByStatus(String status) {
        LambdaQueryWrapper<Img2Img> wrapper = new LambdaQueryWrapper<>();
        wrapper.eq(Img2Img::getStatus, status)
                .orderByDesc(Img2Img::getCreateTime);
        return list(wrapper);
    }

    @Override
    public List<Img2Img> getByModelName(String modelName) {
        LambdaQueryWrapper<Img2Img> wrapper = new LambdaQueryWrapper<>();
        wrapper.eq(Img2Img::getModelName, modelName)
                .orderByDesc(Img2Img::getCreateTime);
        return list(wrapper);
    }

    @Override
    public PageDTO<Img2Img> getPage(int page, int size) {
        Page<Img2Img> pageParam = new Page<>(page, size);
        LambdaQueryWrapper<Img2Img> wrapper = new LambdaQueryWrapper<>();
        wrapper.orderByDesc(Img2Img::getCreateTime);
        Page<Img2Img> result = page(pageParam, wrapper);
        return PageDTO.of(result, Img2Img.class);
    }

    @Override
    public long countRecords() {
        return count();
    }

    @Override
    public long countSuccessRecords() {
        LambdaQueryWrapper<Img2Img> wrapper = new LambdaQueryWrapper<>();
        wrapper.eq(Img2Img::getStatus, "success");
        return count(wrapper);
    }

    @Override
    public long countFailedRecords() {
        LambdaQueryWrapper<Img2Img> wrapper = new LambdaQueryWrapper<>();
        wrapper.eq(Img2Img::getStatus, "failed");
        return count(wrapper);
    }
    
    @Override
    public Map<String, Object> getStats() {
        try {
            Map<String, Object> stats = new HashMap<>();
            stats.put("totalRecords", countRecords());
            stats.put("successRecords", countSuccessRecords());
            stats.put("failedRecords", countFailedRecords());
            
            // 计算成功率
            long total = countRecords();
            if (total > 0) {
                double successRate = (double) countSuccessRecords() / total * 100;
                stats.put("successRate", String.format("%.2f%%", successRate));
            } else {
                stats.put("successRate", "0.00%");
            }
            
            return stats;
        } catch (Exception e) {
            log.error("获取统计信息失败", e);
            throw new RuntimeException("获取统计信息失败: " + e.getMessage(), e);
        }
    }

    @Override
    public String preprocessImage(String module, String inputImage, Integer processorRes,
                                 java.math.BigDecimal thresholdA, java.math.BigDecimal thresholdB) {
        // 此方法已不再使用，Flask API处理预处理功能
        throw new UnsupportedOperationException("此方法已不再使用，Flask API处理预处理功能");
    }
    
    /**
     * 调用Flask API
     */
    private Map<String, Object> callFlaskApi(String endpoint, String base64Image, String prompt, String negativePrompt) {
        try {
            String url = flaskApiUrl + endpoint;
            
            // 准备HTTP请求头
            HttpHeaders headers = new HttpHeaders();
            headers.setContentType(MediaType.MULTIPART_FORM_DATA);
            
            // 将Base64解码为二进制数据
            byte[] imageBytes = Base64.getDecoder().decode(base64Image);
            
            // 构造表单数据
            MultiValueMap<String, Object> body = new LinkedMultiValueMap<>();
            ByteArrayResource imageResource = new ByteArrayResource(imageBytes) {
                @Override
                public String getFilename() {
                    return "image.png";
                }
            };
            body.add("image", imageResource);
            
            // 添加提示词
            if (prompt != null && !prompt.trim().isEmpty()) {
                body.add("prompt", prompt);
            }
            if (negativePrompt != null && !negativePrompt.trim().isEmpty()) {
                body.add("negative_prompt", negativePrompt);
            }
            
            // 创建请求实体
            HttpEntity<MultiValueMap<String, Object>> requestEntity = new HttpEntity<>(body, headers);
            
            // 发送请求
            log.info("发送请求到Flask API: {}", url);
            ResponseEntity<Map> response = restTemplate.exchange(url, HttpMethod.POST, requestEntity, Map.class);
            
            if (response.getStatusCode() == HttpStatus.OK && response.getBody() != null) {
                log.info("Flask API调用成功");
                return response.getBody();
            } else {
                throw new RuntimeException("Flask API调用失败，状态码: " + response.getStatusCode());
            }
            
        } catch (Exception e) {
            log.error("调用Flask API失败", e);
            throw new RuntimeException("调用Flask API失败: " + e.getMessage(), e);
        }
    }
    
    /**
     * 处理API响应
     */
    private Img2ImgResponse processApiResponse(Map<String, Object> apiResponse, Img2ImgRequest request, String modelName) {
        try {
            String taskId = (String) apiResponse.get("task_id");
            String status = (String) apiResponse.get("status");
            
            if (!"success".equals(status)) {
                String errorMessage = apiResponse.containsKey("error_message") ? 
                    (String) apiResponse.get("error_message") : "生成失败，请重试";
                throw new RuntimeException(errorMessage);
            }
            
            // 准备存储所有图片的URL
            List<String> imageUrls = new ArrayList<>();
            String timestamp = String.valueOf(System.currentTimeMillis());
            String dateFolder = java.time.LocalDate.now().toString().replace("-", "/");
            
            // 上传原始图片到OSS
            String originalBase64 = extractBase64Data(request.getInitImage());
            byte[] originalImageData = Base64.getDecoder().decode(originalBase64);
            String originalFileName = String.format("init_images/%s/init_%s_%s.png", dateFolder, taskId, timestamp);
            String initImageUrl = aliOssUtil.upload(originalImageData, originalFileName);
            
            // 处理最终结果图片
            String finalImageBase64 = (String) apiResponse.get("image");
            if (finalImageBase64 != null) {
                byte[] finalImageData = Base64.getDecoder().decode(finalImageBase64);
                String finalFileName = String.format("img2img_images/%s/final_%s_%s.png", dateFolder, taskId, timestamp);
                String finalImageUrl = aliOssUtil.upload(finalImageData, finalFileName);
                imageUrls.add(finalImageUrl);
            }
            
            // 处理Canny边缘检测图片
            String cannyImageBase64 = (String) apiResponse.get("canny_image");
            if (cannyImageBase64 != null) {
                byte[] cannyImageData = Base64.getDecoder().decode(cannyImageBase64);
                String cannyFileName = String.format("img2img_images/%s/canny_%s_%s.png", dateFolder, taskId, timestamp);
                String cannyImageUrl = aliOssUtil.upload(cannyImageData, cannyFileName);
                imageUrls.add(cannyImageUrl);
            }
            
            // 处理OpenPose骨骼图片
            String openposeImageBase64 = (String) apiResponse.get("openpose_image");
            if (openposeImageBase64 != null) {
                byte[] openposeImageData = Base64.getDecoder().decode(openposeImageBase64);
                String openposeFileName = String.format("img2img_images/%s/openpose_%s_%s.png", dateFolder, taskId, timestamp);
                String openposeImageUrl = aliOssUtil.upload(openposeImageData, openposeFileName);
                imageUrls.add(openposeImageUrl);
            }
            
            // 构建响应
            Img2ImgResponse response = new Img2ImgResponse();
            response.setTaskId(taskId);
            response.setStatus("success");
            response.setGenerationTime(apiResponse.get("generation_time") != null ? Long.valueOf(apiResponse.get("generation_time").toString()) : 0L);
            response.setModelName(modelName);
            response.setPrompt((String) apiResponse.get("prompt"));
            response.setNegativePrompt((String) apiResponse.get("negative_prompt"));
            response.setSeed(apiResponse.get("seed") != null ? Long.valueOf(apiResponse.get("seed").toString()) : null);
            response.setSamplerName((String) apiResponse.get("sampler_name"));
            response.setScheduler((String) apiResponse.get("scheduler"));
            response.setSteps(apiResponse.get("steps") != null ? Integer.valueOf(apiResponse.get("steps").toString()) : null);
            response.setCfgScale(apiResponse.get("cfg_scale") != null ? java.math.BigDecimal.valueOf(Double.parseDouble(apiResponse.get("cfg_scale").toString())) : null);
            response.setWidth(apiResponse.get("width") != null ? Integer.valueOf(apiResponse.get("width").toString()) : null);
            response.setHeight(apiResponse.get("height") != null ? Integer.valueOf(apiResponse.get("height").toString()) : null);
            response.setBatchSize(1);
            response.setDenoisingStrength(apiResponse.get("denoising_strength") != null ? java.math.BigDecimal.valueOf(Double.parseDouble(apiResponse.get("denoising_strength").toString())) : null);
            response.setModelHash((String) apiResponse.get("model_hash"));
            response.setRestoreFaces(apiResponse.get("restore_faces") != null ? (Boolean) apiResponse.get("restore_faces") : false);
            response.setTiling(apiResponse.get("tiling") != null ? (Boolean) apiResponse.get("tiling") : false);
            
            // 设置图片URL
            response.setImageUrls(imageUrls);
            response.setInitImageUrl(initImageUrl);
            
            // 保存到数据库
            saveToDatabase(response, request);
            
            return response;
            
        } catch (Exception e) {
            log.error("处理API响应失败", e);
            throw new RuntimeException("处理API响应失败: " + e.getMessage(), e);
        }
    }
    
    /**
     * 保存记录到数据库
     */
    private void saveToDatabase(Img2ImgResponse response, Img2ImgRequest request) {
        try {
            Img2Img record = new Img2Img();
            
            // 设置基本信息
            record.setTaskId(response.getTaskId());
            record.setPrompt(response.getPrompt());
            record.setNegativePrompt(response.getNegativePrompt());
            record.setSeed(response.getSeed());
            record.setSamplerName(response.getSamplerName());
            record.setScheduler(response.getScheduler());
            record.setSteps(response.getSteps());
            record.setCfgScale(response.getCfgScale());
            record.setWidth(response.getWidth());
            record.setHeight(response.getHeight());
            record.setBatchSize(response.getBatchSize());
            record.setDenoisingStrength(response.getDenoisingStrength());
            record.setModelName(response.getModelName());
            record.setModelHash(response.getModelHash());
            record.setRestoreFaces(response.getRestoreFaces());
            record.setTiling(response.getTiling());
            
            // 设置状态信息
            record.setStatus(response.getStatus());
            record.setGenerationTime(response.getGenerationTime());
            
            // 设置图片URL
            record.setInitImageUrl(response.getInitImageUrl());
            record.setImageUrls(objectMapper.writeValueAsString(response.getImageUrls()));
            
            // 设置时间信息
            record.setCreateTime(LocalDateTime.now());
            record.setUpdateTime(LocalDateTime.now());
            
            // 保存到数据库
            save(record);
            
            log.info("成功保存记录到数据库，任务ID: {}", response.getTaskId());
        } catch (Exception e) {
            log.error("保存记录到数据库失败", e);
            throw new RuntimeException("保存记录到数据库失败: " + e.getMessage(), e);
        }
    }
    
    /**
     * 从 data:image/jpeg;base64, 格式中提取纯 base64 数据
     */
    private String extractBase64Data(String dataUrl) {
        if (dataUrl == null || dataUrl.trim().isEmpty()) {
            throw new IllegalArgumentException("图片数据不能为空");
        }

        String processedData = dataUrl.trim();

        // 移除可能的 data URL 前缀
        if (processedData.startsWith("data:")) {
            int commaIndex = processedData.indexOf(",");
            if (commaIndex != -1) {
                processedData = processedData.substring(commaIndex + 1);
            }
        }

        return processedData;
    }
} 